#include "KalmanFilter.h"

KalmanFilter::KalmanFilter(double R, double Q) {
    x = Vector4::Zero();     // Initial state vector [x, y, vx, vy]
    x_ex = Vector4::Zero();  // Extrapolated state

    // Initial state transition matrix F (identity for now)
    F = Matrix4x4::Identity();

    // Measurement matrix H (we only measure x and y position)
    H = Matrix2x4();
    H(0, 0) = 1; H(0, 1) = 0; H(0, 2) = 0; H(0, 3) = 0;
    H(1, 0) = 0; H(1, 1) = 1; H(1, 2) = 0; H(1, 3) = 0;

    // Initial uncertainty covariance P
    P = Matrix4x4::Identity() * 1000;  // High initial uncertainty

    // Measurement noise covariance R (for position)
    this->R = Matrix2x2::Identity() * R;

    // Process noise covariance Q (for position and velocity)
    this->Q = Matrix4x4::Identity() * Q;
}

void KalmanFilter::updateStateTransitionMatrix(double dt) {
    F = Matrix4x4::Identity();
    F(0, 2) = dt;
    F(1, 3) = dt;
}

void KalmanFilter::predict() {
    x = F * x;
    P = F * P * F.transpose() + Q;
}

void KalmanFilter::extrapolate() {
    x_ex = F * x;
}

void KalmanFilter::update(const Vector2& z) {
    // Measurement residual
    Vector2 y = z - (H * x);

    // Kalman gain calculation
    Matrix2x2 S = H * P * H.transpose() + R;
    Matrix4x2 PHt = P * H.transpose();
    Matrix4x2 K = PHt * S.inverse();

    // Update state estimate with measurement
    x = x + K * y;

    // Update covariance estimate
    Matrix4x4 I = Matrix4x4::Identity();
    P = (I - K * H) * P;
}

Vector4 KalmanFilter::getState() const {
    return x;
}

Vector4 KalmanFilter::getStateEx() const {
    return x_ex;
} 